import torch

from utils.compressor import Compressor


class Quantizer(Compressor):
    def __init__(self, bit: int = 8, *args, **kwargs) -> None:
        self.bit = bit
        self.arg = args
        self.kwargs = kwargs

    def __str__(self):
        return f"{self.bit}-bit quantization"

    def encode(self, x: torch.Tensor) -> (float, float):
        with torch.no_grad():
            ma = x.data.max().item()
            mi = x.data.min().item()
            k = ((1 << self.bit) - 1) / (ma - mi)
            b = -mi * k
            x.data = torch.round(k * x.data + b)
        return k, b

    def decode(self, x: torch.Tensor, k: float, b: float) -> None:
        with torch.no_grad():
            x.data -= b
            x.data /= k
